#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <mpi.h>

#define PI 3.14159265358979323846

typedef double complex cplx;

// 串行FFT实现
void fft(cplx *buf, int n) {
    if (n <= 1) return;
    
    cplx *even = malloc(n/2 * sizeof(cplx));
    cplx *odd = malloc(n/2 * sizeof(cplx));
    
    for (int i = 0; i < n/2; i++) {
        even[i] = buf[2*i];
        odd[i] = buf[2*i+1];
    }
    
    fft(even, n/2);
    fft(odd, n/2);
    
    for (int k = 0; k < n/2; k++) {
        cplx t = cexp(-2 * PI * I * k / n) * odd[k];
        buf[k] = even[k] + t;
        buf[k+n/2] = even[k] - t;
    }
    
    free(even);
    free(odd);
}

// 并行FFT主函数
int main(int argc, char **argv) {
    MPI_Init(&argc, &argv);
    
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    if (argc != 2) {
        if (rank == 0) {
            printf("Usage: %s <N>\n", argv[0]);
        }
        MPI_Finalize();
        return 1;
    }
    
    int N = atoi(argv[1]);
    if ((N & (N - 1)) != 0) {
        if (rank == 0) {
            printf("N must be a power of 2\n");
        }
        MPI_Finalize();
        return 1;
    }
    
    // 每个进程处理的数据量
    int local_size = N / size;
    cplx *local_data = malloc(local_size * sizeof(cplx));
    
    // 初始化数据 (仅在rank 0)
    cplx *global_data = NULL;
    if (rank == 0) {
        global_data = malloc(N * sizeof(cplx));
        for (int i = 0; i < N; i++) {
            global_data[i] = sin(2*PI*i/N) + 0.5 * sin(4*PI*i/N);
        }
    }
    
    double start_time = MPI_Wtime();
    
    // 数据分发 (使用MPI_Scatter)
    MPI_Scatter(global_data, local_size, MPI_C_DOUBLE_COMPLEX,
               local_data, local_size, MPI_C_DOUBLE_COMPLEX,
               0, MPI_COMM_WORLD);
    
    // 本地FFT计算
    fft(local_data, local_size);
    
    // 数据收集 (使用MPI_Gather)
    MPI_Gather(local_data, local_size, MPI_C_DOUBLE_COMPLEX,
              global_data, local_size, MPI_C_DOUBLE_COMPLEX,
              0, MPI_COMM_WORLD);
    
    // 全局组合 (仅在rank 0)
    if (rank == 0) {
        // 这里需要添加全局组合的逻辑
        // 由于FFT的特殊性，可能需要额外的通信步骤
    }
    
    double end_time = MPI_Wtime();
    
    if (rank == 0) {
        printf("FFT computation time: %.6f seconds\n", end_time - start_time);
        // 可以添加验证正确性的代码
    }
    
    free(local_data);
    if (rank == 0) {
        free(global_data);
    }
    
    MPI_Finalize();
    return 0;
}